from __future__ import print_function, division
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.autograd import Variable
from resNest import resnest50


print(torch.cuda.is_available())

all_dir = r'E:\samples\dog_cat\for_train'

batch_size = 64
img_size = 224
train_transforms = transforms.Compose([
    transforms.Resize(img_size),
    # transforms.RandomResizedCrop(224),
    transforms.ToTensor(),
    transforms.Normalize((.5, .5, .5), (.5, .5, .5))
])
val_transforms = transforms.Compose([
    transforms.Resize(img_size),
    # transforms.RandomResizedCrop(224),
    transforms.ToTensor(),
    transforms.Normalize((.5, .5, .5), (.5, .5, .5))
])


learning_rate = 0.00005
#动态调整学习率，每30个epoch下降百分之10
def adjust_learning_rate(optimizer, epoch):
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    lr = learning_rate * (0.1 ** (epoch // 300))
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr
        print("learning rate",lr)


def train():
    epochs = 2001
    train_dir = os.path.join(all_dir, "train")
    train_datasets = datasets.ImageFolder(train_dir, transform=train_transforms)
    train_dataloader = torch.utils.data.DataLoader(train_datasets, batch_size=batch_size, shuffle=True)
    cls = train_datasets.classes
    print(cls)
    val_dir = os.path.join(all_dir, "val")
    val_datasets = datasets.ImageFolder(val_dir, transform=val_transforms)
    val_dataloader = torch.utils.data.DataLoader(val_datasets, batch_size=batch_size, shuffle=False)
    print(train_datasets, val_datasets)
    # --------------------训练过程---------------------------------
    net = resnest50(pretrained=False)
    optimizer = torch.optim.Adam(net.fc.parameters(), lr=0.0001)
    print("model", net)
    # state_dict = torch.load(r".\resnest_50_xh.pth")
    # net.load_state_dict(state_dict)
    if torch.cuda.is_available():
        net.cuda()
    # 定义损失函数和优化方式
    criterion = nn.CrossEntropyLoss()  # 损失函数为交叉熵，多用于多分类问题
    optimizer = optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)  # 优化方式为mini-batch momentum-SGD，并采用L2正则化（权重衰减）
    with open("acc.txt", "w") as f:
        with open("log.txt", "w")as f2:
            for epoch in range(0,epochs):
                print('\nEpoch: %d' % (epoch + 1))
                net.train()
                sum_loss = 0.0
                correct = 0.0
                total = 0.0
                adjust_learning_rate(optimizer,epoch)
                for i, data in enumerate(train_dataloader, 0):
                    # 准备数据
                    length = len(train_dataloader)
                    inputs, labels = data
                    inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda()
                    optimizer.zero_grad()
                    import time
                    start = time.perf_counter()
                    outputs = net(inputs)
                    end = time.perf_counter()
                    print('耗时s：', end - start)
                    loss = criterion(outputs, labels)
                    loss.backward()
                    optimizer.step()

                    # 每训练1个batch打印一次loss和准确率
                    sum_loss += loss.item()
                    _, predicted = torch.max(outputs.data, 1)
                    # print(predicted)
                    total += labels.size(0)
                    correct += predicted.eq(labels.data).cpu().sum()
                    print("batch Acc%.3f"%(predicted.eq(labels.data).cpu().sum().item()/batch_size))
                    if epoch % 5 == 0:
                        torch.save(net.state_dict(),'./model/'+ "resNeSt" + str(epoch)+'_dogcat.pth')
                    print('[epoch:[%d/%d], iter:%d/%d] Loss: %.03f | Acc: %.3f%% '
                          % (epoch + 1, epochs, (i + 1), length, sum_loss / (i + 1), 100. * correct / total))
                    f2.write('%03d  %05d |Loss: %.03f | Acc: %.3f%% '
                          % (epoch + 1, (i + 1 + epoch * length), sum_loss / (i + 1), 100. * correct / total))
                    f2.write('\n')
                    f2.flush()

                # 每训练完一个epoch测试一下准确率
                print("Waiting Test!")
                with torch.no_grad():
                    correct = 0
                    total = 0
                    for data in val_dataloader:
                        net.eval()
                        inputs, labels = data
                        inputs, labels = Variable(inputs).cuda(), Variable(labels).cuda()
                        import time
                        start = time.perf_counter()
                        # 当中是你的程序
                        outputs = net(inputs)
                        end = time.perf_counter()
                        print('耗时s：', end - start)
                        # 取得分最高的那个类 (outputs.data的索引号)
                        _, predicted = torch.max(outputs.data, 1)
                        total += labels.size(0)
                        correct += (predicted == labels).sum()
                    print('测试分类准确率为：%.3f%%' % (100 * correct / total))
                    acc = 100. * correct / total
                    # 将每次测试结果实时写入acc.txt文件中
                    print('Saving model......')
                    f.write("EPOCH=%03d,Accuracy= %.3f%%" % (epoch + 1, acc))
                    f.write('\n')
                    f.flush()
                    # 记录最佳测试分类准确率并写入best_acc.txt文件中
                    # if acc > best_acc:
                    #     f3 = open("best_acc.txt", "w")
                    #     f3.write("EPOCH=%d,best_acc= %.3f%%" % (epoch + 1, acc))
                    #     f3.close()
                    #     best_acc = acc
            print("Training Finished, TotalEPOCH=%d" % epoch)


import os
def del_dir(path):
    for i in os.listdir(path):
        path_file = os.path.join(path, i)       # 取文件绝对路径
        if os.path.isfile(path_file):
            os.remove(path_file)
        else:
            del_dir(path_file)


def rm_mkdir(dir_path):
    if os.path.exists(dir_path):
        del_dir(dir_path)
        print('Clean path - %s' % dir_path)
    else:
        os.makedirs(dir_path)
        print('Create path - %s' % dir_path)


if __name__ == "__main__":
    # run()
    train()